package com.project.demo.test;

import it.unisa.dia.gas.jpbc.Element;
import it.unisa.dia.gas.jpbc.Pairing;
import it.unisa.dia.gas.jpbc.PairingParameters;
import it.unisa.dia.gas.plaf.jpbc.pairing.PairingFactory;
import it.unisa.dia.gas.plaf.jpbc.pairing.a.TypeACurveGenerator;

import java.io.*;
import java.util.*;

// 访问树节点类
class AccessTreeNode {
    String attribute;
    int threshold;
    List<AccessTreeNode> children;

    public AccessTreeNode(String attribute, int threshold, List<AccessTreeNode> children) {
        this.attribute = attribute;
        this.threshold = threshold;
        this.children = children;
    }
}

// CP - ABE 加密类
public class CPABEEncryption {

    private Pairing pairing;
    private Element g;
    private Element alpha;
    private Element beta;
    private Element gBeta;
    private Map<String, Element> publicParams;
    private Map<String, Element> masterKey;

    public CPABEEncryption() {
        // 初始化配对参数
        PairingParameters params = new TypeACurveGenerator(512, 512).generate();
        pairing = PairingFactory.getPairing(params);
        g = pairing.getG1().newRandomElement().getImmutable();
        alpha = pairing.getZr().newRandomElement().getImmutable();
        beta = pairing.getZr().newRandomElement().getImmutable();
        gBeta = g.powZn(beta).getImmutable();

        // 生成公共参数
        publicParams = new HashMap<>();
        publicParams.put("g", g);
        publicParams.put("g^beta", gBeta);
        publicParams.put("e(g,g)^alpha", pairing.pairing(g, g).powZn(alpha).getImmutable());

        // 生成主密钥
        masterKey = new HashMap<>();
        masterKey.put("alpha", alpha);
        masterKey.put("beta", beta);
    }

    // 生成用户私钥
    public Map<String, Element> generateUserKey(List<String> attributes) {
        Map<String, Element> userKey = new HashMap<>();
        for (String attribute : attributes) {
            Element D = publicParams.get("g").powZn(masterKey.get("alpha").div(masterKey.get("beta"))).getImmutable();
            userKey.put(attribute, D);
        }
        return userKey;
    }

    // 加密文件，使用访问树作为策略
    public void encryptFile(String inputFilePath, String outputFilePath, AccessTreeNode accessTree) {
        try {
            // 读取文件内容
            byte[] plaintext = readFile(inputFilePath);

            Element m = pairing.getGT().newElementFromBytes(plaintext);
            Element s = pairing.getZr().newRandomElement().getImmutable();
            Element C0 = m.mul(publicParams.get("e(g,g)^alpha").powZn(s)).getImmutable();

            Map<String, byte[]> ciphertextMap = new HashMap<>();
            ciphertextMap.put("C0", C0.toBytes());
            // 递归处理访问树并生成密文组件
            encryptAccessTree(accessTree, s, ciphertextMap);

            // 将密文保存到文件
            try (ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(outputFilePath))) {
                oos.writeObject(ciphertextMap);
            }
            System.out.println("文件加密成功，保存到: " + outputFilePath);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private void encryptAccessTree(AccessTreeNode node, Element s, Map<String, byte[]> ciphertextMap) {
        if (node.children == null || node.children.isEmpty()) {
            // 叶子节点
            Element C = publicParams.get("g").powZn(s).getImmutable();
            ciphertextMap.put(node.attribute, C.toBytes());
        } else {
            // 非叶子节点，进行递归处理
            List<Element> shares = shamirSecretSharing(s, node.threshold, node.children.size());
            for (int i = 0; i < node.children.size(); i++) {
                encryptAccessTree(node.children.get(i), shares.get(i), ciphertextMap);
            }
        }
    }

    private List<Element> shamirSecretSharing(Element secret, int threshold, int numShares) {
        List<Element> shares = new ArrayList<>();
        // 简单示例，实际需要实现 Shamir 秘密共享算法
        for (int i = 0; i < numShares; i++) {
            shares.add(secret);
        }
        return shares;
    }

    // 解密文件，检查用户属性是否满足访问树策略
    public void decryptFile(String inputFilePath, String outputFilePath, Map<String, Element> userKey, AccessTreeNode accessTree) {
        try {
            // 读取加密文件内容
            Map<String, byte[]> ciphertextMap;
            try (ObjectInputStream ois = new ObjectInputStream(new FileInputStream(inputFilePath))) {
                ciphertextMap = (Map<String, byte[]>) ois.readObject();
            }

            Element C0 = pairing.getGT().newElementFromBytes(ciphertextMap.get("C0"));
            if (checkAccessTree(accessTree, userKey, ciphertextMap)) {
                Element m = C0.div(reconstructSecret(accessTree, userKey, ciphertextMap)).getImmutable();
                // 将解密后的内容保存到文件
                writeFile(m.toBytes(), outputFilePath);
                System.out.println("文件解密成功，保存到: " + outputFilePath);
            } else {
                System.out.println("用户属性不满足访问树策略，无法解密文件。");
            }
        } catch (IOException | ClassNotFoundException e) {
            e.printStackTrace();
        }
    }

    private boolean checkAccessTree(AccessTreeNode node, Map<String, Element> userKey, Map<String, byte[]> ciphertextMap) {
        if (node.children == null || node.children.isEmpty()) {
            // 叶子节点
            return userKey.containsKey(node.attribute);
        } else {
            // 非叶子节点，检查子节点满足阈值条件
            int satisfiedCount = 0;
            for (AccessTreeNode child : node.children) {
                if (checkAccessTree(child, userKey, ciphertextMap)) {
                    satisfiedCount++;
                }
            }
            return satisfiedCount >= node.threshold;
        }
    }

    private Element reconstructSecret(AccessTreeNode node, Map<String, Element> userKey, Map<String, byte[]> ciphertextMap) {
        if (node.children == null || node.children.isEmpty()) {
            // 叶子节点
            Element C = pairing.getG1().newElementFromBytes(ciphertextMap.get(node.attribute));
            return pairing.pairing(C, userKey.get(node.attribute));
        } else {
            // 非叶子节点，进行秘密重构
            List<Element> subSecrets = new ArrayList<>();
            for (AccessTreeNode child : node.children) {
                if (checkAccessTree(child, userKey, ciphertextMap)) {
                    subSecrets.add(reconstructSecret(child, userKey, ciphertextMap));
                }
            }
            // 简单示例，实际需要实现秘密重构算法
            Element result = subSecrets.get(0);
            for (int i = 1; i < subSecrets.size(); i++) {
                result = result.mul(subSecrets.get(i));
            }
            return result;
        }
    }

    // 读取文件内容
    private byte[] readFile(String filePath) throws IOException {
        try (FileInputStream fis = new FileInputStream(filePath);
             ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = fis.read(buffer)) != -1) {
                bos.write(buffer, 0, bytesRead);
            }
            return bos.toByteArray();
        }
    }

    // 将解密后的内容保存到文件
    private void writeFile(byte[] content, String filePath) throws IOException {
        try (FileOutputStream fos = new FileOutputStream(filePath)) {
            fos.write(content);
        }
    }

    public static void main(String[] args) {
        CPABEEncryption cpabe = new CPABEEncryption();
        String path = System.getProperty("user.dir")+"/src/main/resources/demo/";

        String inputFilePath =path+"plaintext.txt";
        String encryptedFilePath = path+"encrypted.txt";
        String decryptedFilePath = path+"decrypted.txt";

        // 构建访问树
        AccessTreeNode leaf1 = new AccessTreeNode("ATTR1", 1, null);
        AccessTreeNode leaf2 = new AccessTreeNode("ATTR2", 1, null);
        List<AccessTreeNode> children = Arrays.asList(leaf1, leaf2);
        AccessTreeNode root = new AccessTreeNode(null, 2, children);

        // 生成用户私钥
        List<String> attributes = Arrays.asList("ATTR1", "ATTR2");
        Map<String, Element> userKey = cpabe.generateUserKey(attributes);

        // 加密文件
//        cpabe.encryptFile(inputFilePath, encryptedFilePath, root);

        // 解密文件
        cpabe.decryptFile(encryptedFilePath, decryptedFilePath, userKey, root);
    }
}